import gradio as gr
import numpy as np
from PIL import Image
import logging
import time
from pinecone import Pinecone, ServerlessSpec

# 配置logging
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S"
)

# 初始化Pinecone
class PineconeMNISTPredictor:
    def __init__(self, api_key, index_name="mnist-index"):
        self.pc = Pinecone(api_key=api_key)
        self.index_name = index_name
        self.index = None
        self._connect_to_index()
    
    def _connect_to_index(self):
        """连接到Pinecone索引"""
        try:
            self.index = self.pc.Index(self.index_name)
            logging.info(f"已成功连接到Pinecone索引 '{self.index_name}'")
        except Exception as e:
            logging.error(f"连接到Pinecone索引失败: {e}")
            raise
    
    def preprocess_image(self, img):
        """
        预处理图像：将图像转为28x28的灰度图，
        并将其转换为模型所需的输入格式
        """
        # 将图像转为PIL Image对象（如果不是的话）
        if not isinstance(img, Image.Image):
            img = Image.fromarray(img.astype('uint8'))
        
        # 转为灰度图并调整大小为28x28
        img = img.convert('L')  # 转为灰度图
        img = img.resize((28, 28), Image.Resampling.LANCZOS)  # 调整大小
        
        # 转换为numpy数组并归一化
        img_array = np.array(img)
        img_array = img_array / 255.0  # 归一化到0-1范围
        
        # 反转颜色（因为手写通常是黑色背景白色字，但MNIST是白色背景黑色字）
        img_array = 1.0 - img_array
        
        # 展平为一维数组（784个特征）
        img_flattened = img_array.flatten()
        
        return img_flattened.tolist()
    
    def predict_digit(self, img, k=11):
        """使用Pinecone进行手写数字预测"""
        try:
            # 预处理图像
            processed_img = self.preprocess_image(img)
            
            # 记录开始时间
            start_time = time.time()
            
            # 在Pinecone中查询最近的k个邻居
            query_result = self.index.query(
                vector=processed_img,
                top_k=k,
                include_metadata=True
            )
            
            # 计算推理时间
            inference_time = time.time() - start_time
            
            # 从metadata中获取邻居标签并进行多数投票
            neighbor_labels = []
            for match in query_result['matches']:
                if 'metadata' in match and 'label' in match['metadata']:
                    neighbor_labels.append(match['metadata']['label'])
            
            if not neighbor_labels:
                return "无法找到匹配的邻居", inference_time
            
            # 多数投票决定预测标签
            predicted_label = max(set(neighbor_labels), key=neighbor_labels.count)
            
            # 计算置信度（最高票数/总票数）
            confidence = neighbor_labels.count(predicted_label) / len(neighbor_labels)
            
            return predicted_label, inference_time, confidence
            
        except Exception as e:
            logging.error(f"预测出错: {e}")
            return f"预测出错: {str(e)}", 0, 0

# 初始化Pinecone预测器
# 注意：请替换为你的实际API密钥
API_KEY = "pcsk_5NnFXN_TfqkRHgaAHGhqTL3JHdfe1YbPsaFcZh4JPREN2TBXpRcia1CHqgTUnSqchnDFvx"

try:
    predictor = PineconeMNISTPredictor(api_key=API_KEY)
    logging.info("Pinecone预测器初始化成功")
except Exception as e:
    logging.error(f"Pinecone预测器初始化失败: {e}")
    predictor = None

# 预测函数（供Gradio使用）
def predict_digit_interface(img):
    """Gradio界面的预测函数"""
    if predictor is None:
        return "Pinecone服务未初始化，请检查API密钥和网络连接"
    
    if img is None:
        return "请先绘制一个数字"
    
    result = predictor.predict_digit(img)
    
    if len(result) == 3:
        predicted_digit, inference_time, confidence = result
        return f"预测结果: {predicted_digit}\n置信度: {confidence:.2%}\n推理时间: {inference_time:.2f}秒\n(第一次推理较慢是正常的)"
    else:
        return result[0]

# 创建Gradio界面
with gr.Blocks(title="基于Pinecone的手写数字识别") as demo:
    gr.Markdown("# 基于Pinecone的手写数字识别")
    gr.Markdown("""
    ## 使用Pinecone向量数据库进行实时手写数字识别
    
    在下方的画布上写下一个数字（0-9），然后点击预测按钮查看结果。
    
    **注意**: 第一次推理可能较慢，因为需要初始化网络连接。
    """)
    
    with gr.Row():
        with gr.Column():
            # 创建绘图组件
            sketchpad = gr.Sketchpad(
                height=280, 
                width=280, 
                label="请在此处手写数字",
                brush_radius=8
            )
            
            with gr.Row():
                predict_btn = gr.Button("预测", variant="primary")
                clear_btn = gr.Button("清除")
        
        with gr.Column():
            # 输出结果组件
            output = gr.Textbox(
                label="预测结果",
                lines=4,
                placeholder="预测结果将显示在这里..."
            )
            
            # 添加说明
            gr.Markdown("""
            ### 使用说明:
            1. 在左侧画布上用鼠标绘制一个数字(0-9)
            2. 点击"预测"按钮进行识别
            3. 点击"清除"按钮重新绘制
            
            ### 技术特点:
            - 使用Pinecone向量数据库进行相似度搜索
            - 基于MNIST数据集的k-NN算法
            - 实时推理和结果显示
            """)
    
    # 设置按钮点击事件
    predict_btn.click(
        fn=predict_digit_interface,
        inputs=sketchpad,
        outputs=output
    )
    
    # 清除按钮
    clear_btn.click(
        fn=lambda: None,
        inputs=None,
        outputs=sketchpad,
        queue=False
    )

# 在Jupyter中运行的特殊处理
def launch_in_jupyter():
    """在Jupyter环境中启动应用"""
    try:
        # 检查是否在Jupyter环境中
        if 'get_ipython' in globals():
            print("在Jupyter环境中检测到Gradio应用")
            print("应用将在新窗口中启动...")
            return demo.launch(share=False, inbrowser=True)
        else:
            return demo.launch(share=False)
    except Exception as e:
        print(f"启动应用时出错: {e}")
        return None

# 启动应用
if __name__ == "__main__" or "get_ipython" in globals():
    logging.info("正在启动手写数字识别应用...")
    launch_result = launch_in_jupyter()
    if launch_result:
        logging.info("应用启动成功")
    else:
        logging.error("应用启动失败")